"""
Gradio Web界面
提供友好的用户交互界面
"""

import gradio as gr
import os
import sys
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

from src.models.text2video import TextToVideoGenerator
from src.models.style_transfer import StyleTransfer
from src.models.video_editing import VideoEditor
from src.utils.text_utils import TextProcessor

# 初始化模型
text_processor = TextProcessor()
generator = None
transfer = None
editor = None

def load_models():
    """加载模型"""
    global generator, transfer, editor
    
    try:
        generator = TextToVideoGenerator()
        transfer = StyleTransfer()
        editor = VideoEditor()
        return "✓ 模型加载成功"
    except Exception as e:
        return f"✗ 模型加载失败: {e}"

def generate_video(prompt, style, duration, fps, resolution):
    """生成视频"""
    try:
        if generator is None:
            return None, "请先加载模型"
        
        # 优化提示词
        optimized_prompt = text_processor.optimize_prompt(prompt, style)
        
        # 生成视频
        output_path = generator.generate(
            prompt=optimized_prompt,
            style=style,
            duration=float(duration),
            fps=int(fps),
            resolution=tuple(map(int, resolution.split('x')))
        )
        
        return output_path, f"视频生成成功: {output_path}"
        
    except Exception as e:
        return None, f"生成失败: {e}"

def transfer_style(input_video, style, strength):
    """风格迁移"""
    try:
        if transfer is None:
            return None, "请先加载模型"
        
        # 保存上传的视频
        temp_path = f"data/temp/input_{os.path.basename(input_video.name)}"
        os.makedirs("data/temp", exist_ok=True)
        
        with open(temp_path, "wb") as f:
            f.write(input_video.read())
        
        # 执行风格迁移
        output_path = transfer.transfer(
            input_path=temp_path,
            style=style,
            strength=float(strength)
        )
        
        return output_path, f"风格迁移成功: {output_path}"
        
    except Exception as e:
        return None, f"风格迁移失败: {e}"

def edit_video(input_video, operations, start_time, end_time):
    """编辑视频"""
    try:
        if editor is None:
            return None, "请先加载模型"
        
        # 保存上传的视频
        temp_path = f"data/temp/input_{os.path.basename(input_video.name)}"
        os.makedirs("data/temp", exist_ok=True)
        
        with open(temp_path, "wb") as f:
            f.write(input_video.read())
        
        # 解析操作
        ops = [op.strip() for op in operations.split(',') if op.strip()]
        
        # 执行编辑
        output_path = editor.edit(
            input_path=temp_path,
            operations=ops,
            start_time=float(start_time) if start_time else None,
            end_time=float(end_time) if end_time else None
        )
        
        return output_path, f"视频编辑成功: {output_path}"
        
    except Exception as e:
        return None, f"视频编辑失败: {e}"

def create_gradio_app():
    """创建Gradio应用"""
    
    # 界面标题
    title = "🎬 AI动画生成系统"
    description = """
    基于深度学习的AI动画生成系统，支持文本到动画生成、风格迁移、视频编辑等功能。
    
    ### 功能特性:
    - **文本到动画生成**: 通过自然语言描述生成动画
    - **风格迁移**: 将特定艺术风格应用到视频中
    - **视频编辑**: 智能视频剪辑和特效添加
    
    ### 使用说明:
    1. 点击"加载模型"按钮初始化系统
    2. 选择要使用的功能
    3. 输入参数并点击生成
    4. 等待处理完成并下载结果
    """
    
    with gr.Blocks(title=title, theme=gr.themes.Soft()) as app:
        gr.Markdown(f"# {title}")
        gr.Markdown(description)
        
        # 模型加载
        with gr.Row():
            load_btn = gr.Button("🚀 加载模型", variant="primary")
            model_status = gr.Textbox(label="模型状态", value="未加载", interactive=False)
        
        load_btn.click(load_models, outputs=model_status)
        
        # 功能选项卡
        with gr.Tabs():
            # 文本到动画生成
            with gr.TabItem("🎨 文本到动画生成"):
                with gr.Row():
                    with gr.Column():
                        prompt_input = gr.Textbox(
                            label="文本描述",
                            placeholder="例如：一只可爱的小猫在花园里玩耍",
                            lines=3
                        )
                        style_select = gr.Dropdown(
                            label="生成风格",
                            choices=["realistic", "cartoon", "anime", "oil_painting", 
                                   "watercolor", "sketch", "cyberpunk", "vintage"],
                            value="realistic"
                        )
                        
                        with gr.Row():
                            duration_input = gr.Slider(
                                label="视频时长(秒)",
                                minimum=1.0,
                                maximum=10.0,
                                value=3.0,
                                step=0.5
                            )
                            fps_input = gr.Slider(
                                label="帧率",
                                minimum=8,
                                maximum=30,
                                value=8,
                                step=1
                            )
                        
                        resolution_input = gr.Dropdown(
                            label="分辨率",
                            choices=["256x256", "512x512", "768x768", "1024x1024"],
                            value="512x512"
                        )
                        
                        generate_btn = gr.Button("🎬 生成视频", variant="primary")
                    
                    with gr.Column():
                        video_output = gr.Video(label="生成结果")
                        status_output = gr.Textbox(label="状态", interactive=False)
                
                generate_btn.click(
                    generate_video,
                    inputs=[prompt_input, style_select, duration_input, fps_input, resolution_input],
                    outputs=[video_output, status_output]
                )
            
            # 风格迁移
            with gr.TabItem("🎭 风格迁移"):
                with gr.Row():
                    with gr.Column():
                        style_input_video = gr.Video(label="输入视频")
                        style_select_transfer = gr.Dropdown(
                            label="目标风格",
                            choices=["cartoon", "anime", "oil_painting", "watercolor", 
                                   "sketch", "cyberpunk", "vintage"],
                            value="cartoon"
                        )
                        strength_input = gr.Slider(
                            label="风格强度",
                            minimum=0.1,
                            maximum=1.0,
                            value=0.8,
                            step=0.1
                        )
                        
                        transfer_btn = gr.Button("🎨 风格迁移", variant="primary")
                    
                    with gr.Column():
                        style_output_video = gr.Video(label="风格迁移结果")
                        style_status = gr.Textbox(label="状态", interactive=False)
                
                transfer_btn.click(
                    transfer_style,
                    inputs=[style_input_video, style_select_transfer, strength_input],
                    outputs=[style_output_video, style_status]
                )
            
            # 视频编辑
            with gr.TabItem("✂️ 视频编辑"):
                with gr.Row():
                    with gr.Column():
                        edit_input_video = gr.Video(label="输入视频")
                        operations_input = gr.Textbox(
                            label="编辑操作",
                            placeholder="例如：trim,resize,blur",
                            value="trim"
                        )
                        
                        with gr.Row():
                            start_time_input = gr.Number(
                                label="开始时间(秒)",
                                value=0.0
                            )
                            end_time_input = gr.Number(
                                label="结束时间(秒)",
                                value=5.0
                            )
                        
                        edit_btn = gr.Button("✂️ 编辑视频", variant="primary")
                    
                    with gr.Column():
                        edit_output_video = gr.Video(label="编辑结果")
                        edit_status = gr.Textbox(label="状态", interactive=False)
                
                edit_btn.click(
                    edit_video,
                    inputs=[edit_input_video, operations_input, start_time_input, end_time_input],
                    outputs=[edit_output_video, edit_status]
                )
            
            # 帮助信息
            with gr.TabItem("📚 帮助"):
                help_text = """
                ### 使用指南
                
                **文本到动画生成:**
                - 输入详细的文本描述
                - 选择合适的风格
                - 调整时长、帧率和分辨率
                
                **风格迁移:**
                - 上传要处理的视频
                - 选择目标风格
                - 调整风格强度
                
                **视频编辑:**
                - 上传要编辑的视频
                - 指定编辑操作
                - 设置时间范围
                
                ### 支持的操作
                - trim: 裁剪视频
                - resize: 调整尺寸
                - blur: 添加模糊
                - brightness: 调整亮度
                - contrast: 调整对比度
                - saturation: 调整饱和度
                
                ### 注意事项
                - 首次使用需要加载模型，请耐心等待
                - 生成时间取决于视频长度和复杂度
                - 建议使用GPU加速以获得更好的性能
                """
                gr.Markdown(help_text)
    
    return app

if __name__ == "__main__":
    app = create_gradio_app()
    app.launch(share=True, server_name="0.0.0.0", server_port=7860) 